#!/usr/bin/env python
# -*- coding:utf-8 -*-
__author__ = 'wshu'
__version__ = '1.0'
"""
    ***********************************
    *  @filename : VULNERABLE.py
    *  @Author : wshu
    *  @CodeDate : 2020/6/14 9:02
    *  @Software : PyCharm
    ***********************************
    漏洞模型
"""
from . import BaseModel
from sqlalchemy.sql import expression
from sqlalchemy.orm import relationship
from sqlalchemy import (Column, ForeignKey, Table)
from sqlalchemy.types import (String, Integer, Boolean, Text, DateTime, LargeBinary)

from .ASSETS import Asset
from sqlalchemy_utils.types.choice import ChoiceType
from core.comm.utils import timezone

# 漏洞等级
VULN_LEAVE = (
    ('0', '信息'),
    ('1', '低危'),
    ('2', '中危'),
    ('3', '高危'),
    ('4', '紧急'),
)

# 漏洞状态
VULN_STATUS = (
    ('0', '已忽略'),
    ('1', '已修复'),
    ('2', '待修复'),
    ('3', '漏洞重现'),
    ('4', '复查中'),
)

# 扫描器
SCANNER_TYPE = (
    ('WEB', (
            ('AWVS','AWVS'),
        )
    ),
    ('System', (
               ('Nessus','Nessus'),
               )
    ),
)

class Vulnerable(BaseModel):
    """
    漏洞
    """
    __tablename__ = 'vulnerable'

    cve_id = Column(String(30), comment='漏洞编号')
    cnvd_id = Column(String(30), nullable=True, comment='cnvd编号')
    cve_name = Column(String(240), comment='漏洞名称')
    leave = Column(String(10), comment='危险等级')
    introduce = Column(Text, comment='漏洞简介')
    scopen = Column(Text, comment='影响范围')
    fix = Column(Text, comment='修复方案')
    fix_step = Column(String(150), nullable=True, comment='修复指南')
    update_at = Column(DateTime, default=timezone.now, comment='更新日期')

    def __str__(self):
        return self.cve_id


class Cnvdfiles(BaseModel):
    """
    cnvd
    """
    __tablename__ = 'cnvdfiles'

    title = Column(String(50), comment='文件标题')
    file = Column(LargeBinary)
    update_at = Column(DateTime, default=timezone.now, comment='更新日期')

    def __str__(self):
        return self.title

class VulnScan(BaseModel):
    """
    漏洞扫描
    """
    __tablename__ = 'vulnerable_scan'

    vuln_id = Column(String(30), comment='漏洞编号')
    vuln_name = Column(String(240), comment='漏洞名称')
    cve_num = Column(String(50), nullable=False, comment='cve编号')
    vuln_type = Column(String(60), comment='漏洞属性')
    leave = Column(String(10), ChoiceType(VULN_LEAVE), comment='危险等级')
    introduce = Column(Text, nullable=True, comment='漏洞简介')
    vuln_info = Column(Text, nullable=True, comment='漏洞信息')
    scopen = Column(Text, comment='影响范围')
    fix = Column(Text, nullable=True, comment='修复方案')
    fix_action = Column(Text, nullable=True, comment='处理记录')
    fix_status = Column(String(30), ChoiceType(VULN_STATUS), comment='修复状态')
    create_at = Column(DateTime, default=timezone.now, comment='发现时间')
    update_at = Column(DateTime, default=timezone.now, comment='修复时间')

    vuln_asset = Column(Integer, ForeignKey(f'{Asset.__tablename__}.id', ondelete='CASCADE'))

    def __str__(self):
        return self.vuln_id
